import os
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import transforms, datasets
import torch.cuda.amp as amp
from model import FruitClassifier
from tqdm import tqdm
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix, classification_report
import seaborn as sns
from colorama import Fore, Style
import argparse
import matplotlib.font_manager as fm

# 设置中文字体
try:
    # 尝试设置微软雅黑字体
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei', 'Arial Unicode MS', 'sans-serif']
    plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
except:
    print("警告: 无法设置中文字体，图表中的中文可能无法正确显示")

# 测试数据预处理
test_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

def load_test_data(data_dir, batch_size=32):
    """加载测试数据"""
    test_dataset = datasets.ImageFolder(
        root=data_dir,
        transform=test_transform
    )
    test_loader = DataLoader(
        test_dataset, 
        batch_size=batch_size, 
        shuffle=False, 
        num_workers=2,
        pin_memory=True
    )
    return test_loader, test_dataset

def evaluate_model(model, test_loader, criterion, device, amp_enabled=True):
    """评估模型性能"""
    model.eval()
    test_loss = 0
    test_correct = 0
    test_total = 0
    
    # 存储预测结果和真实标签
    all_preds = []
    all_labels = []
    
    # 存储类别概率
    class_probs = []
    
    with torch.no_grad():
        with amp.autocast(enabled=amp_enabled):
            test_pbar = tqdm(test_loader, desc="测试进度", ncols=100)
            for inputs, labels in test_pbar:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = model(inputs)
                loss = criterion(outputs, labels)
                
                test_loss += loss.item() * inputs.size(0)
                _, predicted = outputs.max(1)
                test_total += labels.size(0)
                test_correct += predicted.eq(labels).sum().item()
                
                # 保存预测和标签
                all_preds.extend(predicted.cpu().numpy())
                all_labels.extend(labels.cpu().numpy())
                
                # 保存类别概率
                probs = torch.nn.functional.softmax(outputs, dim=1)
                class_probs.append(probs.cpu().numpy())
                
                # 更新进度条
                test_acc = 100. * test_correct / test_total
                test_pbar.set_postfix_str(f"L={test_loss/test_total:.4f} acc={test_acc:.2f}%")
    
    test_loss = test_loss / test_total
    test_acc = 100. * test_correct / test_total
    
    # 合并所有批次的类别概率
    class_probs = np.concatenate(class_probs, axis=0)
    
    return {
        "loss": test_loss,
        "accuracy": test_acc,
        "predictions": np.array(all_preds),
        "labels": np.array(all_labels),
        "class_probs": class_probs
    }

def plot_confusion_matrix(y_true, y_pred, classes, save_path=None):
    """绘制混淆矩阵"""
    cm = confusion_matrix(y_true, y_pred)
    plt.figure(figsize=(12, 10))
    
    # 使用英文标签避免中文显示问题
    class_indices = [str(i) for i in range(len(classes))]
    
    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues", xticklabels=class_indices, yticklabels=class_indices)
    plt.xlabel('Predicted Labels')
    plt.ylabel('True Labels')
    plt.title('Confusion Matrix')
    
    # 添加类别映射说明
    plt.figtext(0.5, 0.01, "Class mapping: " + ", ".join([f"{i}: {c}" for i, c in enumerate(classes)]), 
                ha="center", fontsize=8, wrap=True)
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path)
        print(f"混淆矩阵已保存到: {save_path}")
    
    try:
        plt.show()
    except:
        print("无法显示图形，但图形已保存")

def plot_top_misclassified(class_probs, y_true, y_pred, classes, images, n=10, save_path=None):
    """绘制最容易误分类的样本"""
    # 找出误分类的样本
    misclassified = np.where(y_pred != y_true)[0]
    
    if len(misclassified) == 0:
        print("没有误分类的样本！")
        return
    
    # 计算每个误分类样本的置信度
    confidences = np.max(class_probs[misclassified], axis=1)
    
    # 按置信度排序
    sorted_indices = np.argsort(confidences)[::-1]
    
    # 确保索引不会超出范围
    available_images = min(len(images), len(misclassified))
    if available_images == 0:
        print("没有可用的图像来显示误分类样本")
        return
    
    # 只选择有对应图像的误分类样本
    valid_indices = []
    for idx in sorted_indices:
        if misclassified[idx] < len(images):
            valid_indices.append(idx)
            if len(valid_indices) >= n:
                break
    
    if not valid_indices:
        print("没有有效的误分类样本可以显示")
        return
    
    # 确定图表布局
    num_samples = len(valid_indices)
    cols = min(5, num_samples)
    rows = (num_samples + cols - 1) // cols
    
    # 绘制图像
    fig, axes = plt.subplots(rows, cols, figsize=(15, 3*rows))
    if rows * cols > 1:
        axes = axes.flatten()
    else:
        axes = [axes]
    
    for i, idx_in_sorted in enumerate(valid_indices):
        if i < len(axes):
            idx = misclassified[sorted_indices[idx_in_sorted]]
            if idx >= len(images):
                continue
                
            img = images[idx]
            pred_label = classes[y_pred[idx]]
            true_label = classes[y_true[idx]]
            conf = confidences[sorted_indices[idx_in_sorted]]
            
            axes[i].imshow(img)
            # 使用简化的标题避免中文显示问题
            axes[i].set_title(f"Pred: {y_pred[idx]}\nTrue: {y_true[idx]}\nConf: {conf:.2f}")
            axes[i].axis('off')
    
    # 隐藏未使用的子图
    for i in range(len(valid_indices), len(axes)):
        axes[i].axis('off')
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path)
        print(f"误分类图像已保存到: {save_path}")
    
    try:
        plt.show()
    except:
        print("无法显示图形，但图形已保存")

def main():
    parser = argparse.ArgumentParser(description='测试水果分类模型')
    parser.add_argument('--model_path', type=str, default='model/best_model.pth', help='模型路径')
    parser.add_argument('--test_dir', type=str, default='plant_data/Test_Set_Folder', help='测试数据目录')
    parser.add_argument('--batch_size', type=int, default=32, help='批次大小')
    parser.add_argument('--save_results', action='store_true', help='是否保存测试结果')
    parser.add_argument('--output_dir', type=str, default='test_results', help='结果保存目录')
    
    args = parser.parse_args()
    
    # 创建输出目录
    if args.save_results and not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)
    
    # 设置设备
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"{Fore.CYAN}使用设备: {device}{Style.RESET_ALL}")
    
    # 加载测试数据
    test_loader, test_dataset = load_test_data(args.test_dir, args.batch_size)
    classes = test_dataset.classes
    print(f"{Fore.GREEN}加载了 {len(test_dataset)} 个测试样本，共 {len(classes)} 个类别{Style.RESET_ALL}")
    
    # 加载模型
    print(f"{Fore.YELLOW}正在加载模型: {args.model_path}{Style.RESET_ALL}")
    checkpoint = torch.load(args.model_path, map_location=device)
    
    model = FruitClassifier(num_classes=len(classes))
    model.load_state_dict(checkpoint['model_state_dict'])
    model = model.to(device)
    
    # 定义损失函数
    criterion = nn.CrossEntropyLoss()
    
    # 评估模型
    print(f"{Fore.CYAN}开始评估模型...{Style.RESET_ALL}")
    results = evaluate_model(model, test_loader, criterion, device)
    
    # 打印结果
    print(f"\n{Fore.GREEN}测试结果:{Style.RESET_ALL}")
    print(f"测试损失: {results['loss']:.4f}")
    print(f"测试准确率: {results['accuracy']:.2f}%")
    
    # 打印分类报告
    print(f"\n{Fore.GREEN}分类报告:{Style.RESET_ALL}")
    report = classification_report(results['labels'], results['predictions'], target_names=classes)
    print(report)
    
    # 保存分类报告
    if args.save_results:
        with open(os.path.join(args.output_dir, 'classification_report.txt'), 'w') as f:
            f.write(f"测试损失: {results['loss']:.4f}\n")
            f.write(f"测试准确率: {results['accuracy']:.2f}%\n\n")
            f.write(report)
    
    # 绘制混淆矩阵
    print(f"\n{Fore.YELLOW}绘制混淆矩阵...{Style.RESET_ALL}")
    save_path = os.path.join(args.output_dir, 'confusion_matrix.png') if args.save_results else None
    plot_confusion_matrix(results['labels'], results['predictions'], classes, save_path)
    
    # 获取一些测试图像用于可视化
    test_images = []
    max_images = 50  # 限制图像数量以避免内存问题
    image_count = 0
    
    for inputs, _ in test_loader:
        for img in inputs:
            if image_count >= max_images:
                break
                
            img_np = img.permute(1, 2, 0).cpu().numpy()
            # 反归一化
            mean = np.array([0.485, 0.456, 0.406])
            std = np.array([0.229, 0.224, 0.225])
            img_np = std * img_np + mean
            img_np = np.clip(img_np, 0, 1)
            test_images.append(img_np)
            image_count += 1
            
        if image_count >= max_images:
            break
    
    # 绘制误分类样本
    if len(test_images) > 0:
        print(f"\n{Fore.YELLOW}绘制误分类样本...{Style.RESET_ALL}")
        print(f"收集了 {len(test_images)} 张图像用于可视化")
        save_path = os.path.join(args.output_dir, 'misclassified.png') if args.save_results else None
        plot_top_misclassified(
            results['class_probs'][:len(test_images)], 
            results['labels'][:len(test_images)], 
            results['predictions'][:len(test_images)], 
            classes, 
            test_images, 
            n=min(10, len(test_images)),
            save_path=save_path
        )

if __name__ == '__main__':
    main()